from typing import Dict, Any, AsyncGenerator
from agents.components.base import ComponentConfig, ComponentType
from agents.components.llm.ollama_component import OllamaComponent
from agents.components.http.http_component import HttpComponent
from agents.components.vector.chroma_component import ChromaComponent
from agents.flow import Flow, Connection
from config import INDICATORS_COLLECTION
from agents.components.embed.ollama_embedding_component import OllamaEmbeddingComponent

def create_indicator_query_flow() -> Flow:
    """创建指标查询流程"""
    flow = Flow("indicator_query")
    
    # 1. 创建查询解析组件
    query_parser_config = ComponentConfig(
        name="query_parser",
        type=ComponentType.LLM,
        config={
            "system_prompt": "你是一个查询解析器，负责从用户输入中提取指标、维度和过滤条件",
            "user_prompt": "请从以下查询中提取指标、维度和过滤条件：{query}",
            "model_name": "qwen2.5:3b",
            "stream": False,
            "temperature": 0.7,
            "max_tokens": 2048
        },
        input_schema={
            "required": ["query"],
            "properties": {
                "query": {"type": "string"}
            }
        },
        output_schema={
            "required": ["indicators", "dimensions", "filters"],
            "properties": {
                "indicators": {"type": "array"},
                "dimensions": {"type": "array"},
                "filters": {"type": "array"}
            }
        }
    )
    query_parser = OllamaComponent(query_parser_config)
    flow.add_component("query_parser", query_parser)

    embedding_config = ComponentConfig(
        name="embedding",
        type=ComponentType.EMBEDDING,
        config={
            "model_name": "bge-m3:latest"
        },
        input_schema={},
        output_schema={}
    )
    embedding = OllamaEmbeddingComponent(embedding_config)
    flow.add_component("embedding", embedding)
    
    # 2. 创建指标匹配组件
    indicator_matcher_config = ComponentConfig(
        name="indicator_matcher",
        type=ComponentType.VECTOR_DB,
        config={
            "collection_name": INDICATORS_COLLECTION,
            "page_size": 10,
            "n_results": 10,
            "embedding_component": "embedding"
        },
        input_schema={
            "required": ["query_text"],
            "properties": {
                "query_text": {"type": "string"},
                "page": {"type": "integer"},
                "page_size": {"type": "integer"},
                "search_type": {"type": "string"},
                "metadata_filter": {"type": "object"},
                "document_filter": {"type": "object"},
                "n_results": {"type": "integer"}
            }
        },
        output_schema={
            "required": ["records"],
            "properties": {
                "records": {"type": "array"},
                "total": {"type": "integer"},
                "message": {"type": "string"}
            }
        }
    )
    indicator_matcher = ChromaComponent(indicator_matcher_config)
    flow.add_component("indicator_matcher", indicator_matcher)
    
    # 3. 创建维度匹配组件
    dimension_matcher_config = ComponentConfig(
        name="dimension_matcher",
        type=ComponentType.LLM,
        config={
            "system_prompt": "你是一个维度匹配器，负责将用户描述的维度与系统中的维度进行匹配",
            "user_prompt": "请将以下维度描述与系统中的维度进行匹配：{dimensions}",
            "model_name": "gpt-3.5-turbo"
        },
        input_schema={
            "required": ["dimensions"],
            "properties": {
                "dimensions": {"type": "array"}
            }
        },
        output_schema={
            "required": ["matched_dimensions"],
            "properties": {
                "matched_dimensions": {"type": "array"}
            }
        }
    )
    dimension_matcher = OllamaComponent(dimension_matcher_config)
    flow.add_component("dimension_matcher", dimension_matcher)
    
    # 4. 创建过滤条件生成组件
    filter_generator_config = ComponentConfig(
        name="filter_generator",
        type=ComponentType.LLM,
        config={
            "system_prompt": "你是一个过滤条件生成器，负责根据用户描述生成过滤条件",
            "user_prompt": "请根据以下描述生成过滤条件：{filters}",
            "model_name": "gpt-3.5-turbo"
        },
        input_schema={
            "required": ["filters"],
            "properties": {
                "filters": {"type": "array"}
            }
        },
        output_schema={
            "required": ["filter_conditions"],
            "properties": {
                "filter_conditions": {"type": "array"}
            }
        }
    )
    filter_generator = OllamaComponent(filter_generator_config)
    flow.add_component("filter_generator", filter_generator)
    
    # 5. 创建查询参数生成组件
    query_generator_config = ComponentConfig(
        name="query_generator",
        type=ComponentType.LLM,
        config={
            "system_prompt": "你是一个查询参数生成器，负责根据指标、维度和过滤条件生成查询参数",
            "user_prompt": "请根据以下信息生成查询参数：\n指标：{indicators}\n维度：{dimensions}\n过滤条件：{filters}",
            "model_name": "gpt-3.5-turbo"
        },
        input_schema={
            "required": ["indicators", "dimensions", "filters"],
            "properties": {
                "indicators": {"type": "array"},
                "dimensions": {"type": "array"},
                "filters": {"type": "array"}
            }
        },
        output_schema={
            "required": ["query_params"],
            "properties": {
                "query_params": {"type": "object"}
            }
        }
    )
    query_generator = OllamaComponent(query_generator_config)
    flow.add_component("query_generator", query_generator)
    
    # 6. 创建数据查询组件
    data_query_config = ComponentConfig(
        name="data_query",
        type=ComponentType.HTTP,
        config={
            "base_url": "http://114.242.246.245:18095/mdrbase-api/prod/rpt-arc/api",
            "headers": {
                "Content-Type": "application/json"
            },
            "timeout": 30,
            "convert_to_dataframe": True
        },
        input_schema={
            "required": ["query_params"],
            "properties": {
                "query_params": {"type": "object"}
            }
        },
        output_schema={
            "required": ["data"],
            "properties": {
                "data": {"type": "object"}
            }
        }
    )
    data_query = HttpComponent(data_query_config)
    flow.add_component("data_query", data_query)
    
    # 7. 创建图表代码生成组件
    chart_generator_config = ComponentConfig(
        name="chart_generator",
        type=ComponentType.LLM,
        config={
            "system_prompt": "你是一个图表代码生成器，负责根据数据生成合适的图表代码",
            "user_prompt": "请根据以下数据生成图表代码：{data}",
            "model_name": "gpt-3.5-turbo"
        },
        input_schema={
            "required": ["data"],
            "properties": {
                "data": {"type": "object"}
            }
        },
        output_schema={
            "required": ["chart_code"],
            "properties": {
                "chart_code": {"type": "string"}
            }
        }
    )
    chart_generator = OllamaComponent(chart_generator_config)
    flow.add_component("chart_generator", chart_generator)
    
    # 添加连接
    flow.add_connection(Connection(
        source_component="query_parser",
        target_component="indicator_matcher",
        source_output="indicators",
        target_input="indicators"
    ))
    
    flow.add_connection(Connection(
        source_component="query_parser",
        target_component="dimension_matcher",
        source_output="dimensions",
        target_input="dimensions"
    ))
    
    flow.add_connection(Connection(
        source_component="query_parser",
        target_component="filter_generator",
        source_output="filters",
        target_input="filters"
    ))
    
    flow.add_connection(Connection(
        source_component="indicator_matcher",
        target_component="query_generator",
        source_output="matched_indicators",
        target_input="indicators"
    ))
    
    flow.add_connection(Connection(
        source_component="dimension_matcher",
        target_component="query_generator",
        source_output="matched_dimensions",
        target_input="dimensions"
    ))
    
    flow.add_connection(Connection(
        source_component="filter_generator",
        target_component="query_generator",
        source_output="filter_conditions",
        target_input="filters"
    ))
    
    flow.add_connection(Connection(
        source_component="query_generator",
        target_component="data_query",
        source_output="query_params",
        target_input="query_params"
    ))
    
    flow.add_connection(Connection(
        source_component="data_query",
        target_component="chart_generator",
        source_output="data",
        target_input="data"
    ))
    
    return flow

async def run_indicator_query(query: str) -> Dict[str, Any]:
    """运行指标查询流程"""
    flow = create_indicator_query_flow()
    result = await flow.execute({"query": query})
    return result

if __name__ == "__main__":
    import asyncio
    
    async def main():
        # 示例查询
        query = "请帮我查询最近一个月的销售额，按地区分组"
        result = await run_indicator_query(query)
        print("查询结果：", result)
    
    # 运行异步主函数
    asyncio.run(main())